#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文件名: tree.py
功能: 
作者: windows
创建日期: 2025/7/15
版本: 1.0
"""
import os


def generate_filtered_directory_tree(startpath, max_depth=None):
    """
    生成过滤后的目录树结构

    参数:
        startpath: 起始路径
        max_depth: 最大深度 (None表示无限制)
    """
    # 需要显示的文件扩展名
    allowed_extensions = {'.py', '.yaml', '.yml'}

    for root, dirs, files in os.walk(startpath):
        # 排除.venv目录
        if '.venv' in root.split(os.sep):
            continue

        level = root.replace(startpath, '').count(os.sep)
        if max_depth is not None and level > max_depth:
            continue

        # 过滤文件，只保留指定扩展名的文件
        filtered_files = [
            f for f in files
            if os.path.splitext(f)[1].lower() in allowed_extensions
        ]

        # 如果当前目录下没有符合条件的文件和子目录，则不显示
        dirs[:] = [d for d in dirs if d != '.venv']  # 从遍历中排除.venv
        if not filtered_files and not any(dirs):
            continue

        indent = ' ' * 4 * level
        print(f"{indent}{os.path.basename(root)}/")

        subindent = ' ' * 4 * (level + 1)
        for f in filtered_files:
            print(f"{subindent}{f}")


# 示例：生成当前目录的过滤后树形结构，最大深度为3
generate_filtered_directory_tree('.', max_depth=3)